import pickle
from datetime import datetime
from argparse import ArgumentParser
from ibm_model import IBMModel1


def load_sent_pairs(source_file, target_file):
    pairs = []

    source_file = open(source_file, 'r')
    target_file = open(target_file, 'r')

    for src_line, tgt_line in zip(source_file, target_file):
        src_line = [int(i) for i in src_line.strip().split()]
        tgt_line = [int(i) for i in tgt_line.strip().split()]
        pair = (src_line, tgt_line)
        pairs.append(pair)

    source_file.close()
    target_file.close()

    return pairs


def main(option):
    src_vocab = pickle.load(open(option.source_vocab, 'rb'))
    tgt_vocab = pickle.load(open(option.target_vocab, 'rb'))

    pairs = load_sent_pairs(option.source_file, option.target_file)

    lf = len(src_vocab['word2id'])
    le = len(tgt_vocab['word2id'])

    print('lf = ' + str(lf))
    print('le = ' + str(le))

    if option.checkpoint == '':
        model = IBMModel1(le, lf)
    else:
        model = pickle.load(open(option.checkpoint, 'rb'))

    for i in range(option.epochs):
        print('Epoch ' + str(i + 1) + ':\t' + str(datetime.now()))
        model.train(pairs)

    pickle.dump(model, open(option.output_file, 'wb'))


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--source_file', type=str, default='output/fbis.en.10k')
    parser.add_argument('--target_file', type=str, default='output/fbis.zh.10k')
    parser.add_argument('--source_vocab', type=str, default='output/fbis.en.pkl')
    parser.add_argument('--target_vocab', type=str, default='output/fbis.zh.pkl')
    parser.add_argument('--epochs', type=int, default=100)
    parser.add_argument('--output_file', type=str, default='output/model.pkl')
    parser.add_argument('--checkpoint', type=str, default='')
    option = parser.parse_args()

    main(option)
